package link.luyu.plugin.brochain;

import java.math.BigInteger;
import java.util.Arrays;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;
import link.luyu.plugin.brochain.crypto.Hash;
import link.luyu.plugin.brochain.crypto.SM2;
import link.luyu.plugin.brochain.trie.SimpleMerklePatriciaTrie;
import link.luyu.plugin.brochain.types.BlockAndPage;
import link.luyu.plugin.brochain.web3.methods.response.Block;
import link.luyu.plugin.brochain.web3.methods.response.Votes;
import org.apache.tuweni.bytes.Bytes;
import org.apache.tuweni.bytes.Bytes32;
import org.web3j.rlp.RlpEncoder;
import org.web3j.rlp.RlpList;
import org.web3j.rlp.RlpString;
import org.web3j.utils.Numeric;

public class Validator {

    public static boolean validateBlockHash(Block block) {
        byte[] encodedBlockHeader =
                RlpEncoder.encode(
                        new RlpList(
                                RlpString.create(
                                        Numeric.hexStringToByteArray(block.getParentHash())),
                                RlpString.create(
                                        Numeric.hexStringToByteArray(block.getStateRoot())),
                                RlpString.create(
                                        Numeric.hexStringToByteArray(block.getTransactionsRoot())),
                                RlpString.create(
                                        Numeric.hexStringToByteArray(block.getPagesRoot())),
                                RlpString.create(
                                        Numeric.hexStringToByteArray(block.getReceiptsRoot())),
                                RlpString.create(block.getNumber()),
                                RlpString.create(block.getTimestamp()),
                                RlpString.create(
                                        Numeric.hexStringToByteArray(block.getExtraData()))));
        byte[] blockHash = Hash.sm3(encodedBlockHeader);
        return Arrays.equals(blockHash, Numeric.hexStringToByteArray(block.getHash()));
    }

    public static boolean validateVotes(Votes votes, Set<String> consensusNodes) {
        if (consensusNodes == null || consensusNodes.size() == 0) {
            return false;
        }

        if (votes.getBlockNumber().equals(BigInteger.ZERO)) {
            return true;
        }

        byte[] params = Numeric.hexStringToByteArray(votes.getParams());
        byte[] blockNumber = Numeric.toBytesPadded(votes.getBlockNumber(), 8);
        byte[] blockHash = Numeric.hexStringToByteArray(votes.getBlockHash());

        byte[] message = new byte[109];
        System.arraycopy(params, 0, message, 0, 1);
        System.arraycopy(blockNumber, 0, message, 1, 8);
        System.arraycopy(params, 1, message, 9, 4);
        System.arraycopy(blockHash, 0, message, 13, 32);

        Set<String> votedNodes = new HashSet<>();
        for (Votes.Vote vote : votes.getVotes()) {
            if (!consensusNodes.contains(vote.getNodeId())) {
                continue;
            }

            byte[] nodeId = Numeric.hexStringToByteArray(vote.getNodeId());
            byte[] signature = Numeric.hexStringToByteArray(vote.getSignature());
            BigInteger r = Numeric.toBigInt(signature, 0, 32);
            BigInteger s = Numeric.toBigInt(signature, 32, 32);
            byte v = signature[64];
            System.arraycopy(nodeId, 0, message, 45, 64);

            BigInteger publicKey = SM2.recoverFromSignature(v, r, s, message, true);

            if (publicKey == null) {
                return false;
            }
            if (!Arrays.equals(Numeric.toBytesPadded(publicKey, 64), nodeId)) {
                return false;
            }

            votedNodes.add(vote.getNodeId());
        }

        return votedNodes.size() * 3 >= consensusNodes.size() * 2;
    }

    public static boolean validateBlockPagesRoot(Block block) {
        List<byte[]> encodedRegionPages =
                block.getRegionPages()
                        .stream()
                        .map(
                                regionPage ->
                                        new RlpList(
                                                RlpString.create(
                                                        Numeric.hexStringToByteArray(
                                                                regionPage.getRegionId())),
                                                RlpString.create(
                                                        Numeric.hexStringToByteArray(
                                                                regionPage.getPageHash()))))
                        .map(RlpEncoder::encode)
                        .collect(Collectors.toList());

        SimpleMerklePatriciaTrie<Bytes, byte[]> trie = new SimpleMerklePatriciaTrie<>(Bytes::wrap);
        for (int i = 0; i < encodedRegionPages.size(); i++) {
            Bytes key = Bytes.wrap(RlpEncoder.encode(RlpString.create(BigInteger.valueOf(i))));
            trie.put(key, encodedRegionPages.get(i));
        }
        Bytes32 rootHash = trie.getRootHash();

        return rootHash.equals(Bytes32.wrap(Numeric.hexStringToByteArray(block.getPagesRoot())));
    }

    public static boolean validateBlockAndPage(
            BlockAndPage blockAndPage, Set<String> consensusNodes) {
        return Validator.validateBlockHash(blockAndPage.getBlock())
                && Validator.validateBlockPagesRoot(blockAndPage.getBlock())
                && Validator.validateVotes(blockAndPage.getVotes(), consensusNodes);
    }
}
